package com.example.mobilenet

import android.graphics.Bitmap
import android.graphics.BitmapFactory
import android.graphics.Matrix
import android.os.Bundle
import android.util.Log
import android.widget.Button
import android.widget.CheckBox
import android.widget.ImageView
import android.widget.TextView
import androidx.appcompat.app.AppCompatActivity
import com.example.mobilenet.databinding.ActivityMainBinding
import java.io.InputStream
import java.util.AbstractMap
import java.util.Arrays
import kotlin.math.exp

class MainActivity : AppCompatActivity() {
    val TEST_IMAGE=  "panda.jpeg"

    companion object {
        const val TAG = "TNet_MainActivity"

        const val MNN_MODEL = "mobilenet_v2.caffe.mnn"//"mobilenet_v2_1.0_224.tflite.mnn"
        const val MNN_LIB = "libMNN_dl.so"

        init {
            System.loadLibrary("mobilenet")
        }
    }

    private lateinit var binding: ActivityMainBinding
    private lateinit var btnRun : Button
    private lateinit var fileUtil: FileUtil
    private lateinit var imageOrigin :ImageView
    private lateinit var imageClassified :ImageView
    private lateinit var tvClassified :TextView
    private var isRunning = false
    private var frameWorkType = InterpreterJni.kOnnxRuntime

    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        fileUtil = FileUtil(this)
        binding = ActivityMainBinding.inflate(layoutInflater)
        setContentView(binding.root)

        frameWorkType = InterpreterJni.kMnn

        imageOrigin = findViewById(R.id.iv_origin)
        imageClassified = findViewById(R.id.iv_classified)
        tvClassified = findViewById(R.id.tv_classified)
        val inputStream: InputStream = assets.open(TEST_IMAGE)
        val bmpTest = BitmapFactory.decodeStream(inputStream)
        imageOrigin.setImageBitmap(bmpTest)

        val tvTips : TextView = findViewById(R.id.tv_tips)
        btnRun = findViewById(R.id.btn_test)
        btnRun.setOnClickListener {
            if (!isRunning) {
                btnRun.text="Running"
                btnRun.isEnabled = false
                tvTips.text = "正在推理..."
                imageClassified.setImageBitmap(null)
            }

            Thread {
                classifyImage(frameWorkType, bmpTest)
                runOnUiThread{
                    btnRun.text="RUN"
                    btnRun.isEnabled = true
                    tvTips.text = ""
                }
            }.start()
        }
    }

    //触发图像分类识别
    private fun classifyImage(frmType : Int, picBitmap: Bitmap?) {
        Log.i(TAG, "classifyImage engine frameWorkType: ${InterpreterJni.getEngineString(frmType)}")
        if (picBitmap == null) {
            Log.i(TAG, "fail picBitmap == null")
            return
        }
        val(libPath, modelPath) =  getLibAndModelPath(frmType)
        if (libPath.isEmpty() || modelPath.isEmpty()) {
            Log.i(TAG, "fail libPath: $libPath, modelPath: $modelPath")
            return
        }
        val inst = InterpreterJni.crateInterpreter(frmType, libPath, modelPath)
        if (inst == 0L) {
            onStepError("crateInterpreter fail inst == 0L")
            return
        }

        //ResNet模型输入 [1,3,224,224]
        val inputTensorShape = InterpreterJni.getInputDataShape(inst)
        Log.i(TAG, "inputTensorShape: ${Arrays.toString(inputTensorShape)}")

        //读取图片并resize,获取color数组
        val height = inputTensorShape[2]
        val width = inputTensorShape[3]

        //ResNet模型输出 [1,1000]
        val outputTensorShape = InterpreterJni.getOutputDataShape(inst)
        Log.i(TAG, "outputTensorShape: ${Arrays.toString(outputTensorShape)}")
        val length = outputTensorShape[1]
        //[1000]
        val resultArray = FloatArray(length)

        var inputPixelsArray = intArrayOf()
        inputPixelsArray = getResizedImageColor(height, width, picBitmap) //[224 * 224]

        if (inputPixelsArray.isEmpty()) {
            onStepError("classifyImage inputPixelsArray isEmpty!")
            return
        }
        //通过JNI触发Native底层模型运行
        val runResult = InterpreterJni.run(inst, inputPixelsArray, resultArray)
        InterpreterJni.destroyInterpreter(inst)
        if (runResult != "success") {
            onStepError("classifyImage InterpreterJni.run fail: $runResult")
            return
        }
        Log.i(TAG, "classifyImage InterpreterJni.run success !!!, resultArray index 0 value: ${resultArray[0]}")
        var processModelResult = false

        //结果后处理
        processModelResult = processModelResult(resultArray)

        if (!processModelResult) {
            onStepError("classifyImage fail, not processModelResult")
        } else {
            Log.i(TAG, "classifyImage success done~")
        }
    }

    //处理模型推理结果
    private fun processModelResult(resultArray: FloatArray): Boolean {

        //如果是tflite模型，需去除第 0 类（background），从第 1 类开始处理
//        val subArray = resultArray.sliceArray(1 until resultArray.size)

        val softMax = resultArray.mapIndexed { index, value -> ClassifyLabel(
            LabelMap.Labels[index], value) }
            .sortedByDescending { it.confidence }
        var classifyResult = "识别结果(前5):\n"
        Log.i(TAG, "识别结果(前5):")
        for (index in 0 until 5) {
            val label = softMax[index]
            Log.i(TAG, "${index + 1}. $label")
            classifyResult += "${index + 1}. $label\n"
        }
        tvClassified.text = classifyResult
        tvClassified.setTextColor(resources.getColor(R.color.black))
        return true
    }

    private fun onStepError(error: String) {
        runOnUiThread {
            Log.e(TAG, error)
            tvClassified.text = error
            tvClassified.setTextColor(resources.getColor(R.color.red))
        }
    }

    //获取指定推理引擎所需的SO库文件地址
    private fun getLibAndModelPath(frmType : Int) : Pair<String, String> {
        var libPath = fileUtil.getLocalPathByAssetName(MNN_LIB)
        var modelPath = fileUtil.getLocalPathByAssetName(MNN_MODEL)
        return Pair(libPath, modelPath)
    }

    //获取图像的color数组，数组每个Int值代表一个像素点color(argb) //hongry_test5 得到的是argb
    private fun getResizedImageColor(h: Int, w: Int, picBitmap: Bitmap) : IntArray {
        val resizedBmp = resizeImage(picBitmap, w, h)
        val pixels = IntArray(resizedBmp.width * resizedBmp.height)
        resizedBmp.getPixels(pixels, 0, resizedBmp.width, 0, 0, resizedBmp.width, resizedBmp.height)

        return pixels
    }

    //将图片调整为目标所需尺寸
    private fun resizeImage(inputImage: Bitmap, newWidth: Int, newHeight: Int): Bitmap {
        val matrix = Matrix()
        val scaleWidth = newWidth.toFloat() / inputImage.width.toFloat()
        val scaleHeight = newHeight.toFloat() / inputImage.height.toFloat()
        matrix.postScale(scaleWidth, scaleHeight)
        return Bitmap.createBitmap(
            inputImage, 0, 0, inputImage.width, inputImage.height, matrix, true
        )
    }
}